// SPDX-License-Identifier: GPL-2.0 OR BSD-3-Clause
/* Copyright (c) 2020 Facebook */
#include "vmlinux.h"
#include <bpf/bpf_helpers.h>
#include <bpf/bpf_tracing.h>
#include <bpf/bpf_core_read.h>

char LICENSE[] SEC("license") = "Dual BSD/GPL";


struct {
	__uint(type, BPF_MAP_TYPE_LPM_TRIE);
	__uint(max_entries, 16);
	__type(key, u64);
	__type(value, u32);
} blacklist SEC(".maps");

SEC("xdp")
int firewall(struct xdp_md *ctx) {
	void *data_end = (void *)(long)ctx->data_end;
	void *data = (void *)(long)ctx->data;

	// Only IPv4 supported for this example
	struct ethhdr *ether = data;
	if (data + sizeof(*ether) > data_end) {
		// Malformed Ethernet header
		return XDP_ABORTED;
	}

	if (ether->h_proto != 0x08U) {  // htons(ETH_P_IP) -> 0x08U
		// Non IPv4 traffic
		return XDP_PASS;
	}

	data += sizeof(*ether);
	struct iphdr *ip = data;
	if (data + sizeof(*ip) > data_end) {
		// Malformed IPv4 header
		return XDP_ABORTED;
	}

	struct {
		__u32 prefixlen;
		__u32 saddr;
	} key;

	key.prefixlen = 32;
	key.saddr = ip->saddr;

	// Lookup SRC IP in blacklisted IPs
	__u64 *rule_idx = bpf_map_lookup_elem(&blacklist, &key);
	if (rule_idx) {
		// Matched, increase match counter for matched "rule"
		__u32 index = *(__u32*)rule_idx;  // make verifier happy
		return XDP_DROP;
	}

	return XDP_PASS;
}